import pymysql
import logging
import time
from typing import List, Tuple
from db_config import DBConfig  # 导入统一的DBConfig

# 导入脚本1中的函数
from stock_financial import (
    sync_stock_finance_to_db,
    batch_sync_stock_finance
)

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class StockFinanceTrigger:
    def __init__(self):
        self.db_config = DBConfig.config  # 使用统一的DBConfig
        self.conn = None
        self.connect_db()
    
    def connect_db(self):
        """连接数据库"""
        try:
            self.conn = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                charset=self.db_config['charset']
            )
            logger.info("数据库连接成功")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise
    
    def get_stock_codes_with_names(self, query_sql: str) -> List[Tuple[str, str]]:
        """
        根据查询SQL获取股票代码和名称列表
        
        :param query_sql: 查询SQL语句，需要返回stock_code和stock_name字段
        :return: 股票代码和名称列表 [(code, name), ...]
        """
        stocks = []
        try:
            with self.conn.cursor() as cursor:
                cursor.execute(query_sql)
                results = cursor.fetchall()
                
                for row in results:
                    if len(row) >= 2:
                        stock_code = row[0]  # stock_code字段
                        stock_name = row[1]  # stock_name字段
                        
                        # 根据股票代码添加交易所后缀
                        if stock_code.startswith('6'):
                            formatted_code = f"{stock_code}.SH"
                        elif stock_code.startswith('0') or stock_code.startswith('3') :
                            formatted_code = f"{stock_code}.SZ"
                        elif stock_code.startswith('8') or stock_code.startswith('4') or stock_code.startswith('9'):
                            formatted_code = f"{stock_code}.BJ"
                        else:
                            logger.warning(f"无法判断交易所的股票代码: {stock_code}")
                            continue
                            
                        stocks.append((formatted_code, stock_name))
                    else:
                        # 如果SQL只返回stock_code，尝试自动判断交易所并使用代码作为名称
                        stock_code = row[0]
                        if stock_code.startswith('6'):
                            formatted_code = f"{stock_code}.SH"
                        elif stock_code.startswith('0') or stock_code.startswith('3') or stock_code.startswith('8') or stock_code.startswith('9') or stock_code.startswith('4'):
                            formatted_code = f"{stock_code}.SZ"
                        else:
                            logger.warning(f"无法判断交易所的股票代码: {stock_code}")
                            continue
                            
                        stocks.append((formatted_code, stock_code))
                
                logger.info(f"查询到 {len(stocks)} 只股票")
                return stocks
                
        except Exception as e:
            logger.error(f"执行查询SQL失败: {e}")
            return []
    
    def trigger_stock_finance_sync(self, query_sql: str, batch_size: int = 10, delay_seconds: int = 2):
        """
        触发股票财务数据同步
        
        :param query_sql: 查询SQL语句，用于获取需要同步的股票代码和名称
        :param batch_size: 每批处理的股票数量
        :param delay_seconds: 每批股票同步后的延迟时间（秒）
        """
        # 获取股票代码和名称列表
        stock_list = self.get_stock_codes_with_names(query_sql)
        
        if not stock_list:
            logger.warning("没有找到需要同步的股票")
            return
        
        logger.info(f"开始同步 {len(stock_list)} 只股票的财务数据")
        
        # 分批处理
        for i in range(0, len(stock_list), batch_size):
            batch = stock_list[i:i + batch_size]
            logger.info(f"正在处理第 {i//batch_size + 1}/{(len(stock_list)-1)//batch_size + 1} 批，本批 {len(batch)} 只股票")
            
            # 调用脚本1中的批量同步函数
            batch_sync_stock_finance(batch)
            
            # 如果不是最后一批，添加延迟
            if i + batch_size < len(stock_list):
                logger.info(f"等待 {delay_seconds} 秒后继续下一批...")
                time.sleep(delay_seconds)
        
        logger.info("所有股票财务数据同步完成")
    
    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
            logger.info("数据库连接已关闭")

# 使用示例
if __name__ == "__main__":
    trigger = StockFinanceTrigger()
    
    try:
        # 示例1：同步所有非ST股票
        query_sql = """
        SELECT 
            A.stock_code, 
            A.stock_name 
        FROM 
            stock_basic AS A
        WHERE 
            A.is_stock_st = 0
            AND A.stock_code NOT IN (
                SELECT DISTINCT(stock_code) FROM finance_metadata
            )
        """
        
        # 示例2：同步特定行业的股票
        # query_sql = """
        # SELECT stock_code, stock_name 
        # FROM stock_basic 
        # WHERE industry_type LIKE '%银行%' 
        # AND is_stock_st = 0
        # """
        
        # 示例3：同步特定市场的股票
        # query_sql = """
        # SELECT stock_code, stock_name 
        # FROM stock_basic 
        # WHERE market_type = 'sh_main' 
        # AND is_stock_st = 0
        # """
        
        # 触发同步
        trigger.trigger_stock_finance_sync(
            query_sql=query_sql,
            batch_size=10,          # 每批处理5只
            delay_seconds=0.01         # 每批同步后延迟2秒
        )
        
    except Exception as e:
        logger.error(f"触发同步过程中发生错误: {e}")
    
    finally:
        trigger.close()